import numpy as np

#This script dipslays the power spectrum of the system when called. 

def file_len(filename):
        with open(filename) as f:
                for i, l in enumerate(f):
                        pass
        return i + 1

def get_kolmogorov(len, coeff):
        k_kolm = np.linspace(0, len, num=100)
        p_kolm = coeff * (k_kolm**(-5.0/3.0))
        return k_kolm, p_kolm

def read_spectra(filename):
	#Read as spectrum file written with write_spectra
        nlines = file_len(filename)
        nlines = nlines - 1 #Do not count th headers 
        print "Lines in " + filename + " = " + str(nlines)

        file = open(filename, 'r')

        line = file.readline()
        line = line.split()
        Ptot = line[1]
        line = file.readline()

        vktot_norm = np.zeros(nlines,dtype=float)
        vkx_norm = np.zeros(nlines,dtype=float)
        vky_norm = np.zeros(nlines,dtype=float)
        vkz_norm = np.zeros(nlines,dtype=float)
        kx = np.zeros(nlines,dtype=float)
        ky = np.zeros(nlines,dtype=float)
        kz = np.zeros(nlines,dtype=float)

        ii=0
        for line in file:
                line = line.split()
                vktot_norm[ii] = line[0]
                vkx_norm[ii] = line[1]
                vky_norm[ii] = line[2]
                vkz_norm[ii] = line[3]
                kx[ii] = line[4]
                ky[ii] = line[5]
                kz[ii] = line[6]
                ii = ii + 1

        file.close()

        return Ptot, vktot_norm, vkx_norm, vky_norm, vkz_norm, kx, ky, kz

def write_spectra(filename, vktot_norm, vkx_norm, vky_norm, vkz_norm, kx,  ky,  kz, Ptot):
	#Write powerspectrum generated by normalize_powerspectrum into a file
        file = open(filename, 'w')
        file.write("Ptot %e\n" % Ptot)
        file.write("  vktot_norm     vkx_norm     vky_norm     vkz_norm           kx           ky           kz\n")
        for ii in range(vktot_norm.size):
                file.write("%e %e %e %e %e %e %e\n" % (vktot_norm[ii], vkx_norm[ii], vky_norm[ii], vkz_norm[ii], kx[ii],  ky[ii],  kz[ii]))
        file.close()

def powerspectrum(grid, nx0, ny0, nz0, Lbox):
	nx = nx0 - 6 #Take out the ghost zones 
	ny = ny0 - 6
	nz = nz0 - 6
	imaxx = nx/2
	imaxy = ny/2
	imaxz = nz/2

	# Calculate the k-axis for the power spectra
	#kx = np.linspace(0.0, nx/(2.0*Lbox), num=imaxx)
	kx = 2.0*np.pi*np.fft.fftfreq(nx, d=Lbox/nx)
	kx = kx[0:imaxx]
	#ky = np.linspace(0.0, ny/(2.0*Lbox), num=imaxy)
	ky = 2.0*np.pi*np.fft.fftfreq(ny, d=Lbox/ny)
	ky = ky[0:imaxy]
	#kz = np.linspace(0.0, nz/(2.0*Lbox), num=imaxz)
	kz = 2.0*np.pi*np.fft.fftfreq(nz, d=Lbox/nz)
	kz = kz[0:imaxz]

	# Averaged 1-D spectrum in x-direction
	vkx = np.zeros(imaxx)
	for n in range(0,nz):
		for m in range(0,ny): 
			AA = np.fft.fft(grid[:,m,n])
			for j in range(0,imaxx):
				vkx[j] = vkx[j] + 2.0*np.abs(AA[j])
	vkx = vkx/(ny*nz) 

	# Averaged 1-D spectrum in y-direction
	vky = np.zeros(imaxy)
	for n in range(0,nz):
		for l in range(0,nx): 
			AA = np.fft.fft(grid[l,:,n])
			for j in range(0,imaxy):
				vky[j] = vky[j] + 2.0*np.abs(AA[j])
	vky = vky/(nx*nz)

	# Averaged 1-D spectrum in z-direction
	vkz = np.zeros(imaxz)
	for m in range(0,ny):
		for l in range(0,nx): 
			AA = np.fft.fft(grid[l,m,:])
			for j in range(0,imaxz):
				vkz[j] = vkz[j] + 2.0*np.abs(AA[j])
	vkz = vkz/(nx*ny)
	
	# TODO: Figure out 3-D shell-integrated spectrum

	return vkx, vky, vkz, kx, ky, kz 

def normalize_powerspectrum(uu_x, uu_y, uu_z, COMP_DOMAIN_SIZE_X, COMP_DOMAIN_SIZE_Y, COMP_DOMAIN_SIZE_Z, Lbox):
        #Gather the powerpectra and normalize it with total power
        vkx_x,  vky_x,  vkz_x,  kx,  ky,  kz  = powerspectrum(uu_x,  COMP_DOMAIN_SIZE_X, COMP_DOMAIN_SIZE_Y, COMP_DOMAIN_SIZE_Z, Lbox)
        vkx_y,  vky_y,  vkz_y,  kx,  ky,  kz  = powerspectrum(uu_y,  COMP_DOMAIN_SIZE_X, COMP_DOMAIN_SIZE_Y, COMP_DOMAIN_SIZE_Z, Lbox)
        vkx_z,  vky_z,  vkz_z,  kx,  ky,  kz  = powerspectrum(uu_z,  COMP_DOMAIN_SIZE_X, COMP_DOMAIN_SIZE_Y, COMP_DOMAIN_SIZE_Z, Lbox)

        #Sum all velocity components
        vkx = vkx_x + vkx_y + vkx_z
        vky = vky_x + vky_y + vky_z
        vkz = vkz_x + vkz_y + vkz_z

        #Normalize the powerspectrum with total power
        Ptot = np.sum(vkx) + np.sum(vky) + np.sum(vkz)
        vktot = vkx + vky + vkz
        vkx_norm = vkx/Ptot
        vky_norm = vky/Ptot
        vkz_norm = vkz/Ptot
        vktot_norm = vktot/Ptot

        return Ptot, vktot_norm, vkx_norm, vky_norm, vkz_norm, kx,  ky,  kz


'''
;;
;; $Id: power_snapshot.pro 18246 2012-02-14 10:35:36Z anders@astro.lu.se $
;;
;; Calculate energy spectrum of 3-D cube.
;;
;;   var           : 3-D variable to make power spectrum of
;;   eks           : shell-integrated spectrum
;;   vkx           : 1-D averaged spectrum <f(kx,y,z)>_(y,z)
;;   vky           : 1-D averaged spectrum <f(x,ky,z)>_(x,z)
;;   vkz           : 1-D averaged spectrum <f(x,y,kz)>_(x,y)
;;   ks            : scalar shell wavenumber
;;   kx, ky, kz    : scalar wavenumber in x, y and z
;;   k0x, k0y, k0z : unit wavenumbers
;;   Lx, Ly, Lz    : box dimensions, gives unit wavenumbers k0=2*!pi/L
;;   nks           : number of shells to sum over
;;   deltak        : radius difference of shells
;;
pro power_snapshot, var, eks=eks, vkx=vkx, vky=vky, vkz=vkz, $
    nshell=nshell, ks=ks, kx=kx, ky=ky, kz=kz, $
    k0x=k0x, k0y=k0y, k0z=k0z, Lx=Lx, Ly=Ly, Lz=Lz, deltak=deltak, nks=nks, $
    double=double, plot=plot, ps=ps, filename=filename, $
    nolegend=nolegend, quiet=quiet, debug=debug
;
;  Default values.
;
default, plot, 0
default, ps, 0
default, nolegend, 0
default, quiet, 0
default, debug, 0
;
zero=0.0
one =1.0
if (keyword_set(double)) then begin
  zero=0.0d
  one =1.0d
  var=double(var)
endif
;
sizevar=size(var)
;
if (sizevar[0] ge 1) then nx=sizevar[1] else nx=1
if (sizevar[0] ge 2) then ny=sizevar[2] else ny=1
if (sizevar[0] ge 3) then nz=sizevar[3] else nz=1
;
;  1-D spectrum in x-direction
;
if (arg_present(vkx)) then begin
  vkx=fltarr(nx/2)*zero
  for n=0,nz-1 do begin & for m=0,ny-1 do begin
    AA=fft(reform(var[*,m,n]),-1)
    for j=0,nx/2-1 do vkx[j] = vkx[j] + 2*abs(AA[j])
  endfor & endfor
  vkx=vkx/(ny*nz)
endif
;
;  1-D spectrum in y-direction
;
if (arg_present(vky)) then begin
  vky=fltarr(ny/2)*zero
  for n=0,nz-1 do begin & for l=0,nx-1 do begin
    AA=fft(reform(var[l,*,n]),-1)
    for j=0,ny/2-1 do vky[j] = vky[j] + 2*abs(AA[j])
  endfor & endfor
  vky=vky/(nx*nz)
endif
;
;  1-D spectrum in z-direction
;
if (arg_present(vkz)) then begin
  vkz=fltarr(nz/2)*zero
  for m=0,ny-1 do begin & for l=0,nx-1 do begin
    AA=fft(reform(var[l,m,*]),-1)
    for j=0,nz/2-1 do vkz[j] = vkz[j] + 2*abs(AA[j])
  endfor & endfor
  vkz=vkz/(nx*ny)
endif
;
;  3-D shell-integrated spectrum.
;
if (arg_present(eks)) then begin
;
;  Define directional wavenumbers kx, ky, kz.
;
  if (n_elements(k0x) eq 0) then k0x=one else k0x=k0x*one
  if (n_elements(k0y) eq 0) then k0y=one else k0y=k0y*one
  if (n_elements(k0z) eq 0) then k0z=one else k0z=k0z*one
  if (n_elements(Lx) ne 0)  then k0x=2*!pi/Lx*one
  if (n_elements(Ly) ne 0)  then k0y=2*!pi/Ly*one
  if (n_elements(Lz) ne 0)  then k0z=2*!pi/Lz*one
  if (nx eq 1) then kx=[0] else kx=[indgen(nx/2+1),-reverse(indgen(nx/2-1)+1)]*k0x
  if (ny eq 1) then ky=[0] else ky=[indgen(ny/2+1),-reverse(indgen(ny/2-1)+1)]*k0y
  if (nz eq 1) then kz=[0] else kz=[indgen(nz/2+1),-reverse(indgen(nz/2-1)+1)]*k0z
;
;  Define scalar wavenumber over which to bin.
;
  default, deltak, 1.0*one
  if (n_elements(nks) eq 0) then begin
    ks=indgen(nx/2)*one
  endif else begin
    ks=indgen(nks)*deltak
  endelse
  if (not quiet) then begin
    print, 'k0x, k0y, k0z=', k0x, k0y, k0z
    print, 'Going to sum in shells of radius ks=', ks
  endif
;
;  Define array to hold shell summed power and number of elements in each shell.
;  
  eks=fltarr(n_elements(ks))
  nshell=lonarr(n_elements(ks))
;
;  Transform to wavenumber space.
;
  vkk=fft(var)
;
;  Loop over all vector k, calculate |k| and assign to proper shell.
;
  for ikz=0,nz-1 do begin & for iky=0,ny-1 do begin & for ikx=0,nx-1 do begin
    k=round(sqrt(kx[ikx]^2+ky[iky]^2+kz[ikz]^2)/deltak)
    if (debug) then print, ikx, kx[ikx], iky, ky[iky], ikz, kz[ikz], k
    if (k lt n_elements(ks)) then begin
      eks[k]=eks[k]+abs(vkk[ikx,iky,ikz])^2
      nshell[k]=nshell[k]+1
    endif else begin
      if (debug) then print, '** Mode not allocated to any shell'
    endelse
  endfor & endfor & endfor
;
;  Print total number of modes that were allocated to a shell.
;
  if (not quiet) then begin
    print, 'total(nshell)=', total(nshell), ' / total(modes)=', n_elements(vkk)
  endif
endif
;
;  Make simple plot if requested.
;
if (plot) then begin
  pmin=min(vkx[1:nx/2-1])
  pmax=max(vkx[1:nx/2-1])
  if (ny gt 1) then begin
    pmin=min([pmin,vky[1:ny/2-1]])
    pmax=max([pmax,vky[1:ny/2-1]])
  endif
  if (nz gt 1) then begin
    pmin=min([pmin,vkz[1:nz/2-1]])
    pmax=max([pmax,vkz[1:nz/2-1]])
  endif

  linestyles=[0,1,2]

  if (ps) then begin
    default, filename, 'power.eps'
    set_plot, 'ps'
    device, /encapsulated, color=color, xsize=8.7, ysize=8.0, $
        font_size=11, filename=filename
    !p.font=-1
    !p.charsize=1.0
    thick=3
    !p.charthick=thick & !p.thick=thick & !x.thick=thick & !y.thick=thick
  endif
  
  plot, vkx, xtitle='k/k0', ytitle='|g(k)|', $
      xrange=[1.0,nx/2], $
      yrange=[10.0^floor(alog10(pmin)),10.0^ceil(alog10(pmax))], $
      /xlog, /ylog, $
      linestyle=linestyles[0]
  if (ny gt 1) then oplot, vky, linestyle=linestyles[1]
  if (nz gt 1) then oplot, vkz, linestyle=linestyles[2]
;  legend, ['k!Dx!N','k!Dy!N','k!Dz!N'], linestyle=linestyles, /bottom
  if (not nolegend) then legend, ['1','2','3'], linestyle=linestyles, /bottom

  if (ps) then begin
    device, /close
    set_plot, 'x'
  endif
endif

'''
